import os
import torch
from dataclasses import dataclass
from tqdm import tqdm

from torch.utils.data import DataLoader
from dino_mix.dataset.satgeoloc import SatGeoLocDatasetEval
from dino_mix.transforms.transforms import get_transforms_val
from dino_mix.evaluate.satgeoloc import evaluate, calc_sim
from dino_mix.model import TimmModel
from pathlib import Path
import sys
from dino_mix.utils import Logger
import shutil


@dataclass
class Configuration:
    # Model
    model = 'dinov2_vitb14_MixVPR'

    # backbone
    backbone_arch = 'dinov2_vitb14'
    # backbone_arch = ''
    pretrained = True
    # layers_to_freeze = 1
    # layers_to_crop = []
    layer1 = 7
    use_cls = True
    norm_descs = True

    # Aggregator 聚合方法
    agg_arch = 'MixVPR'  # CosPlace, NetVLAD, GeM
    agg_config = {'in_channels': 768,
                  'in_h': 32,  # 受输入图像尺寸的影响
                  'in_w': 32,
                  'out_channels': 1024,
                  'mix_depth': 2,
                  'mlp_ratio': 1,
                  'out_rows': 4}

    # Override model image size
    img_size: int = 448
    new_hight = 448
    new_width = 448

    # Evaluation
    batch_size: int = 40
    verbose: bool = True
    gpu_ids: tuple = (0, 1)
    normalize_features: bool = True
    neighbour_range = 64

    # save_top_k_img = False
    # save_top_k_img_path = r'none'
    # save_top_k_img_num = 5

    # Dataset
    data_folder = r'E:\datasets\satgeoloc_dataset'

    # Checkpoint to start from
    checkpoint_start = r'weights_e5_98.9714.pth'
    # set num_workers to 0 if on Windows
    num_workers: int = 4

    # train on GPU if available
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'


# -----------------------------------------------------------------------------#
# Config                                                                      #
# -----------------------------------------------------------------------------#

config = Configuration()

if __name__ == '__main__':

    # -----------------------------------------------------------------------------#
    # Model                                                                       #
    # -----------------------------------------------------------------------------#

    file_name, ext = os.path.splitext(os.path.basename(config.checkpoint_start))
    path = Path(config.checkpoint_start)
    save_path = path.parent.absolute()

    sys.stdout = Logger(os.path.join(save_path, 'cvcities_' + file_name + '_test_eval_log.txt'))

    print("\nModel: {}".format(config.model))

    model = TimmModel(model_name=config.model,
                      pretrained=True,
                      img_size=config.img_size, backbone_arch=config.backbone_arch, agg_arch=config.agg_arch,
                      agg_config=config.agg_config, layer1=config.layer1)

    print(model)
    data_config = model.get_config()
    print(data_config)
    mean = data_config["mean"]
    std = data_config["std"]
    img_size = config.img_size

    image_size_sat = (img_size, img_size)

    # new_width = config.img_size * 2
    # new_hight = config.img_size

    new_width = config.new_width
    new_hight = config.new_hight
    img_size_ground = (new_hight, new_width)

    # load pretrained Checkpoint    
    if config.checkpoint_start is not None:
        print("Start from:", config.checkpoint_start)
        model_state_dict = torch.load(config.checkpoint_start)
        model.model.load_state_dict(model_state_dict, strict=False)

        # Data parallel
    print("GPUs available:", torch.cuda.device_count())
    if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
        model = model.to(config.device)
        model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
        print("Model: DataParallel")

    # Model to device   
    # model = model.to(config.device)

    print("\nImage Size Sat:", image_size_sat)
    print("Image Size Ground:", img_size_ground)
    print("Mean: {}".format(mean))
    print("Std:  {}\n".format(std))

    # -----------------------------------------------------------------------------#
    # DataLoader                                                                  #
    # -----------------------------------------------------------------------------#
    # Eval transforms
    sat_transforms_val, ground_transforms_val = get_transforms_val(image_size_sat,
                                                                   img_size_ground,
                                                                   mean=mean,
                                                                   std=std,
                                                                   )
    # reference Images Test
    reference_dataset_test = SatGeoLocDatasetEval(data_folder=config.data_folder,
                                             split="test",
                                             img_type="query",
                                             transforms=ground_transforms_val,

                                             )

    reference_dataloader_test = DataLoader(reference_dataset_test,
                                       batch_size=config.batch_size,
                                       num_workers=config.num_workers,
                                       shuffle=False,
                                       pin_memory=True)

    # Query Images Test
    query_dataset_test = SatGeoLocDatasetEval(data_folder=config.data_folder,
                                             split="test",
                                             img_type="query",
                                             transforms=ground_transforms_val,
                                             )

    query_dataloader_test = DataLoader(query_dataset_test,
                                       batch_size=config.batch_size,
                                       num_workers=config.num_workers,
                                       shuffle=False,
                                       pin_memory=True)


    # -----------------------------------------------------------------------------#
    # Evaluate                                                                    #
    # -----------------------------------------------------------------------------#

    print("Reference Images Test:", len(reference_dataset_test))
    print("Query Images Test:", len(query_dataset_test))

    r1_test = evaluate(config=config,
                       model=model,
                       reference_dataloader=reference_dataloader_test,
                       query_dataloader=query_dataloader_test,
                       ranks=[1, 5, 10],
                       step_size=1000,
                       cleanup=True)
